# -*- coding: utf-8 -*-
"""
Created on Wed Nov 18 09:22:00 2020
策略为三种：
一，直接选保真度最高的
二，如果保真度不比之前的高，就选次优的
三，如果保真度不比之前的高，就选最差的
两个比特各自可以选择不同的动作，但每个比特上，选用的动作都是单方向的。
@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
import copy

np.random.seed(1)
T = 10*np.pi
dt = np.pi/4
step_max = T/dt

action_space = np.array([[1,0,0], #可以选择的动作范围，各列的每项分别代表着 sigma x, y, z 前面的系数。
                       [2,0,0], #每次执行的动作都是单独的绕 x, y, z 轴一定角度的旋转
                       [0,1,0], # x, y 方向的值可以取负，但 z 方向的只能取正值
                       [0,2,0],
                       [0,0,1],
                       [0,0,2],
                       [-1,0,0],
                       [-2,0,0],
                       [0,-1,0],
                       [0,-2,0],
                       [0,0,0]])

s_x = np.mat([[0,1],[1,0]],dtype=complex)
s_y = np.mat([[0,-1j],[1j,0]],dtype=complex)
s_z = np.mat([[1,0],[0,-1]],dtype=complex)
H_c = np.kron((s_x+1j*s_y)/2, (s_x-1j*s_y)/2) + np.kron((s_x-1j*s_y)/2, (s_x+1j*s_y)/2)

init_set_num = 512
target_set_num = init_set_num

qubit_1 = np.mat([[1,0],[0,0]],dtype=complex)
qubit_2 = np.mat([[0,0],[0,1]],dtype=complex)



alpha_num = 8
theta = [np.pi/8,np.pi/4,3*np.pi/8]
theta_1 = theta
theta_2 = theta
theta_3 = theta

alpha = np.linspace(0,np.pi*2,alpha_num,endpoint=False)
alpha_1 = alpha
alpha_2 = alpha
alpha_3 = alpha
alpha_4 = alpha

a_list_complex = np.matrix([[0,0,0,0]],dtype=complex) #第一行用来占位，否则无法和其他行并在一起，在最后要注意去掉这一行
for ii in range(3): #theta_1
    for jj in range(3): #theta_2
        for kk in range(3): #theta_3
            for mm in range(alpha_num): #alpha_1
                for nn in range(alpha_num): #alpha_2
                    for oo in range(alpha_num): #alpha_3
                        for pp in range(alpha_num): #alpha_4
                            
                            a_1_mo = np.cos(theta_1[ii])
                            a_2_mo = np.sin(theta_1[ii])*np.cos(theta_2[jj])
                            a_3_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.cos(theta_3[kk])
                            a_4_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.sin(theta_3[kk])
                            
                            a_1_real = a_1_mo*np.cos(alpha_1[mm])
                            a_1_imag = a_1_mo*np.sin(alpha_1[mm])
                            a_2_real = a_2_mo*np.cos(alpha_2[nn])
                            a_2_imag = a_2_mo*np.sin(alpha_2[nn])
                            a_3_real = a_3_mo*np.cos(alpha_3[oo])
                            a_3_imag = a_3_mo*np.sin(alpha_3[oo])
                            a_4_real = a_4_mo*np.cos(alpha_4[pp])
                            a_4_imag = a_4_mo*np.sin(alpha_4[pp])
                            
                            a_1_complex = a_1_real + a_1_imag*1j
                            a_2_complex = a_2_real + a_2_imag*1j
                            a_3_complex = a_3_real + a_3_imag*1j
                            a_4_complex = a_4_real + a_4_imag*1j
                            
                            a_complex = np.matrix([[ a_1_complex, a_2_complex, a_3_complex, a_4_complex]])
                            a_list_complex = np.row_stack((a_list_complex,a_complex))
                            
psi_set = np.array(np.delete(a_list_complex,0,axis=0)) # 删除矩阵的第一行
np.random.shuffle(psi_set) #打乱顺序
init_set = psi_set[:init_set_num]


def target_set():
    target_set = psi_set[init_set_num : init_set_num + target_set_num]
    return target_set

#动作直接选最优的
def step0(psi,target_psi,F):
    fid_list = []
    psi_list = []
    
    for action1 in range(len(action_space)): #action1 是第一个比特上采用的动作
        for action2 in range(len(action_space)): #action2 是第二个比特上采用的动作
        
            H_12 = (np.kron(qubit_1,action_space[action1,0]*s_x) + np.kron(qubit_1,action_space[action1,1]*s_y) + np.kron(qubit_1,action_space[action1,2]*s_z)+ np.kron(qubit_2,action_space[action2,0]*s_x) + np.kron(qubit_2,action_space[action2,1]*s_y) + np.kron(qubit_2,action_space[action2,2]*s_z))/2
    
            H = H_12 + H_c
            U = expm(-1j * H * dt) 
            psi_ = U * psi
            fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
            
            psi_list.append(psi_)
            fid_list.append(fid)
    best_action = fid_list.index(max(fid_list))
    best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_


#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    
    for action1 in range(len(action_space)): #action1 是第一个比特上采用的动作
        for action2 in range(len(action_space)): #action2 是第二个比特上采用的动作
        
            H_12 = (np.kron(qubit_1,action_space[action1,0]*s_x) + np.kron(qubit_1,action_space[action1,1]*s_y) + np.kron(qubit_1,action_space[action1,2]*s_z)+np.kron(qubit_2,action_space[action2,0]*s_x) + np.kron(qubit_2,action_space[action2,1]*s_y) + np.kron(qubit_2,action_space[action2,2]*s_z))/2
    
            H = H_12 + H_c
            U = expm(-1j * H * dt) 
            psi_ = U * psi
            fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
            
            psi_list.append(psi_)
            fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    for action1 in range(len(action_space)): #action1 是第一个比特上采用的动作
        for action2 in range(len(action_space)): #action2 是第二个比特上采用的动作
        
            H_12 = (np.kron(qubit_1,action_space[action1,0]*s_x) + np.kron(qubit_1,action_space[action1,1]*s_y) + np.kron(qubit_1,action_space[action1,2]*s_z)+np.kron(qubit_2,action_space[action2,0]*s_x) + np.kron(qubit_2,action_space[action2,1]*s_y) + np.kron(qubit_2,action_space[action2,2]*s_z))/2
    
            H = H_12 + H_c
            U = expm(-1j * H * dt) 
            psi_ = U * psi
            fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
            
            psi_list.append(psi_)
            fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        psi_list_ = copy.deepcopy(psi_list)
        fid_list_ = copy.deepcopy(fid_list)
        
        del psi_list_[fid_list_.index(max(fid_list_))]
        del fid_list_[fid_list_.index(max(fid_list_))]
        
        best_action = fid_list.index(max(fid_list_))
        
        best_fid = max(fid_list_)
        
    psi_ = psi_list[best_action]
    
    return best_action, best_fid, psi_
#---------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
def job(target_psi):
    fids_list = []
    target_psi = np.mat(target_psi).T
    for k in range(init_set_num):
        psi1 = np.mat(init_set[k]).T
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        
        fid_max = F
        fid_max1 = F
        fid_max2 = F
        fid_max0 = F
        
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step0(psi,target_psi,F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max1,fid_max2,fid_max0)  
        fids_list.append(fid_max)
        
    return  np.mean(fids_list)


def multicore():
    pool = mp.Pool()
    F_list = pool.map(job, target_set)
    return F_list

    
if __name__ == '__main__':
    target_set = target_set()
    time1 = time()
    F_list = multicore()
    print(F_list)
    print(np.mean(F_list))
    time2 = time()
    print('time_cost is: ',time2-time1)

[0.9604630936138029, 0.9478697960182404, 0.9890370914664826, 0.974271903422741, 0.9560800803662872, 0.9833701590853647, 0.9523085985254923, 0.9664661451412464, 0.9748864515945057, 0.9780919605789054, 0.97664780360974, 0.9713377184046632, 0.9907254499150646, 0.9644908405497623, 0.9627013496557061, 0.9862286143946145, 0.9853589032409892, 0.9652359015366654, 0.9781925428411216, 0.9815543583480553, 0.966056120348084, 0.98191942487134, 0.9348557993872062, 0.967680728167089, 0.9554954418154351, 0.9734749059989254, 0.9734076705955312, 0.9564273662737055, 0.9887341244760723, 0.9560874263584656, 0.9307442567049331, 0.9877163976521248, 0.9893517146417864, 0.9664545814961003, 0.9787334927464616, 0.984931126794794, 0.9400073711673487, 0.971920768406593, 0.9809657384969787, 0.9279979119554889, 0.9573767260799834, 0.9851707619488776, 0.9442930820379184, 0.9748959912171887, 0.9702387506028505, 0.983540828244588, 0.9687745303007758, 0.978762644256693, 0.9653419691163834, 0.94766311082838, 0.9615369507936924, 0.9819605689585353, 0.9888084812351469, 0.9617341267153708, 0.9718362832383154, 0.9432438543929151, 0.9696343438698773, 0.9737638388468978, 0.973118978725836, 0.9744847831328924, 0.9619910209580469, 0.9535602521834776, 0.953068075216054, 0.9762454510960399, 0.9875579042356307, 0.9855585495843637, 0.9574251501146026, 0.98541600569928, 0.9731775201955336, 0.9899355586048599, 0.9836808994279739, 0.9755547786049121, 0.9706041944928432, 0.9751807308048135, 0.9296774101507654, 0.9690143095843755, 0.9429347693509351, 0.9601208225102891, 0.9844156192123237, 0.9570564120515745, 0.9646306076551561, 0.9691466042876751, 0.9873591592218851, 0.9311275927834936, 0.9861077807951732, 0.9578078873912659, 0.9621367328605628, 0.9650540828484855, 0.9891258457531786, 0.9884208745853527, 0.9775384121012264, 0.9878302815613083, 0.9545047821492376, 0.9515101854840708, 0.9610524155743819, 0.9400730419324832, 0.9608872977541744, 0.9604124338822818, 0.9753352819812062, 0.9580194660320953, 0.9323093266177301, 0.9901840215231705, 0.9832694859186479, 0.9707283907319806, 0.9730685117843209, 0.9872926640204983, 0.962309381991619, 0.9588173802712696, 0.9612366003999211, 0.9452262263436542, 0.9680500283963991, 0.9533326156929525, 0.9881192622723772, 0.9884190747647532, 0.9699661471990819, 0.9546852610233783, 0.9834381567155456, 0.9828866187039682, 0.9707417412444289, 0.9691721516299858, 0.9466724748537787, 0.9706409816117523, 0.9694136393353845, 0.9647399510761577, 0.9847460372069532, 0.9640144415373982, 0.9568019803501472, 0.9793312589370238, 0.9823557576275931, 0.9367168016145964, 0.9626170096959766, 0.9885490180134643, 0.9694401451570772, 0.9882284146516228, 0.9877407404331662, 0.938958078957052, 0.9871271192347726, 0.9872271425201999, 0.979271286018556, 0.9765458952882081, 0.9887491471211854, 0.9824288071690712, 0.9730583889976064, 0.9884766008249193, 0.9959584651185829, 0.9707797877287956, 0.9892222823789234, 0.9827133510249524, 0.970052776877643, 0.9632617861530356, 0.9756976335783183, 0.987362177170301, 0.9712767322346374, 0.9607505748123343, 0.9545661488757388, 0.98743359514552, 0.9857144142809644, 0.9888401301242753, 0.9890307318715638, 0.9731055314127166, 0.9675712474490805, 0.9737227430746613, 0.9892375232618128, 0.9887936254926357, 0.9745591604308669, 0.9411533269138871, 0.9672058660779425, 0.945533006721981, 0.9908525148834479, 0.9726199128213731, 0.974036981512065, 0.9721928831767845, 0.9891691199427316, 0.9433665137787514, 0.9844113618571402, 0.9688850302501189, 0.965300028784042, 0.990322635309267, 0.9905450912437562, 0.9715351025661908, 0.972344708892938, 0.9820725949084737, 0.9681130379611582, 0.9875212581062514, 0.9895625063422169, 0.9892787426916609, 0.8943893387095053, 0.9815740522387804, 0.9654337838357105, 0.9795328609550358, 0.9479821697863215, 0.987742567888954, 0.9791089060875311, 0.9745077755175221, 0.9573862698070328, 0.9714957970805573, 0.9697671942883286, 0.9565669778031085, 0.9680475031688558, 0.9673035438221664, 0.9888203039612541, 0.9788725211700651, 0.981407443137595, 0.9358621032475436, 0.9641236259365533, 0.968051934911996, 0.9895625063422169, 0.9903469013066785, 0.9807752254270294, 0.9850159518984919, 0.9625367648093035, 0.9882628893215303, 0.9657372718359563, 0.9879418677102383, 0.9861447397377595, 0.9641996049059701, 0.9659818488858456, 0.9933114782703032, 0.9515177223492749, 0.964347154079438, 0.9296435089796102, 0.9725162790113557, 0.9639179417201467, 0.9920656265790095, 0.9901679915405625, 0.9903593880685267, 0.932620590404103, 0.9762015591470508, 0.9782889895416275, 0.9829611903889679, 0.9820388392306567, 0.9891258457531787, 0.9392165307702816, 0.9527211333177277, 0.9856326153582143, 0.9884006583246197, 0.9631856276154465, 0.9653419691163834, 0.9659979439390187, 0.9935013350041667, 0.9726702834105616, 0.9778725691168229, 0.9747494476732229, 0.9538000411968877, 0.9848772863893525, 0.9748617251582796, 0.9563932229721028, 0.9705240774479537, 0.979824462874746, 0.9795912927546058, 0.9630719454063558, 0.9776129895143146, 0.9737449206325057, 0.9801822153028044, 0.9641938559416993, 0.9960699764807435, 0.9571142377361853, 0.9889245556410011, 0.9689406722632226, 0.9642465861632836, 0.9566818351721287, 0.9689400894263707, 0.977238854509715, 0.991570982002417, 0.9686470164772223, 0.9794172080725594, 0.9414158170973501, 0.9507051760227251, 0.9502474659804601, 0.9600695330889417, 0.9579050581483, 0.9712690249962538, 0.9673732731783669, 0.9879841040470372, 0.9696411775033846, 0.9409934733589415, 0.9687878892920218, 0.9483398335703714, 0.9661340591991576, 0.9746451569988908, 0.9612083408773486, 0.9556224921927532, 0.967977552378597, 0.9789346138154539, 0.9510156793230822, 0.9849241298227386, 0.9707797877287956, 0.9868166042672554, 0.9671552598381999, 0.9895392522028821, 0.9837295243227716, 0.9857468838558705, 0.9877515064179552, 0.953917788322953, 0.9420706340343727, 0.9804623708445291, 0.9881252452485578, 0.984483952347719, 0.9839472412926151, 0.9589314896951808, 0.9650540828484854, 0.9648203489130847, 0.9551570451513299, 0.939556010348183, 0.9805511905303503, 0.9726306996982764, 0.9881115639456524, 0.9664280585586007, 0.9619186222632535, 0.9888600026941003, 0.9891577945984937, 0.9889833135209765, 0.9467469518751768, 0.9883170838562685, 0.9883871267016779, 0.9633516020184286, 0.9469049959420033, 0.9861991876957341, 0.973282891967872, 0.9736454775489274, 0.9719518424493403, 0.9604124338822818, 0.9899045111551945, 0.9892082219306301, 0.9704181274557524, 0.9623256242201494, 0.9767685973257199, 0.9880865507019552, 0.9839057303171617, 0.9824740439936632, 0.9816475169905088, 0.9582157912115336, 0.9752657248128882, 0.9683996704102198, 0.9793479458639156, 0.9759744253410094, 0.9778725671521272, 0.9472237830680893, 0.9903793652074993, 0.9881519771997355, 0.9633511882585911, 0.93411831079883, 0.9720855501059381, 0.9764659037306109, 0.971000581153346, 0.9742425246796895, 0.9644571097065857, 0.963606959846631, 0.9735414422753405, 0.9656654327833927, 0.9851147385853327, 0.942661282230012, 0.9596416033071755, 0.9615953321888056, 0.990832914992066, 0.9613327526298799, 0.9721196169142498, 0.9648659497600471, 0.9736890548832885, 0.9610465781995949, 0.9661293686478569, 0.9697510150572953, 0.9769452081195323, 0.9858436844992414, 0.9897750080727916, 0.9792222019846202, 0.9757248582833351, 0.9670510773214998, 0.9622976080647637, 0.9827753253622475, 0.9881230908141911, 0.9710433747242875, 0.9687644263653694, 0.9828903035432406, 0.9748197188210541, 0.9824980500063927, 0.960596839973973, 0.9872739198528312, 0.9564356858101752, 0.9752369629793383, 0.9707427286937114, 0.9855309935437581, 0.9788507011533967, 0.9448318244180363, 0.9883539458359063, 0.9433048527475036, 0.9625559207729026, 0.9466324577663493, 0.9472207809218669, 0.9772603656176301, 0.9753341830776572, 0.9798083953706012, 0.9691285973236767, 0.9695595076482493, 0.980195195550327, 0.956886599174285, 0.9908792027130544, 0.9674477167337386, 0.9786973188768686, 0.9520660523573634, 0.9645500759408848, 0.9364717623480125, 0.982326398436975, 0.9835032554297887, 0.9759124655835696, 0.9491662918634753, 0.976271167876875, 0.9629704639614415, 0.96293308451545, 0.9451562194746228, 0.9467887232793457, 0.9805282064208471, 0.9894039568482293, 0.9670195566894499, 0.988015893900509, 0.9764397103906288, 0.9638757424833619, 0.9914431423500428, 0.9765039160732893, 0.9691721516299858, 0.97420033700341, 0.9527211333177277, 0.9872728967268597, 0.9821479695407609, 0.9684253533529965, 0.9888881106362251, 0.9468186603429811, 0.9602581201969604, 0.9723566446460383, 0.9790780631898804, 0.8896464346307642, 0.970092249453545, 0.9837171640703654, 0.9679541919171348, 0.9672907861195568, 0.9796563286451866, 0.9898850656752991, 0.9480270891874718, 0.9856366613223011, 0.9650763857034084, 0.9617034650757605, 0.9809292818422117, 0.9880859378433616, 0.9759211453948583, 0.9448181047756479, 0.9693014370150019, 0.9863940564003522, 0.9738690771776073, 0.9747833944531976, 0.99382665919997, 0.9375746424817069, 0.9761352526444631, 0.964361441357955, 0.9275658170513754, 0.9659297044230963, 0.9572917625308334, 0.9861843499206087, 0.9840372192725991, 0.9700787664724755, 0.9736842693805374, 0.965942061938555, 0.9649995127588136, 0.9863988681138083, 0.9916567866815575, 0.9654337838357105, 0.9749956703862918, 0.968960118750859, 0.9719800834497547, 0.9444770265615875, 0.9840693113261261, 0.9841325056005545, 0.9268370638909411, 0.9848311528113297, 0.9674850017825062, 0.9671622193228886, 0.9665845293871037, 0.9622332865864025, 0.976147086194389, 0.9751973533050591, 0.9801222538389456, 0.9570564120515745, 0.9885717847380653, 0.9749956703862919, 0.9699721216785304, 0.9766829571866585, 0.9650305390005791, 0.9840274398346143, 0.9723476340075521, 0.9653813253605377, 0.9739336316969051, 0.9781307134957811, 0.9802897359051143, 0.9948401082043035, 0.9734749059989252, 0.9858827252426358, 0.959429090862443, 0.9460548670921669, 0.982929495603547, 0.9459354674773355, 0.9825932696910685, 0.9473915420882466, 0.9556016910350791, 0.9797538219897308, 0.945148187381612, 0.9786962451168177, 0.9867884045999689, 0.9888052779829033, 0.9790823633412808, 0.8988041900245363, 0.9780060429936357, 0.9893266973278113, 0.9560261736280293]
0.9706291110784504
time_cost is:  179719.48897218704

0.00-0.49:  0
0.50-0.54:  0 0 0 0 0
0.55-0.59:  0 0 0 0 0
0.60-0.64:  0 0 0 0 0
0.65-0.69:  0 0 0 0 0
0.70-0.74:  0 0 0 0 0
0.75-0.79:  0 0 0 0 0
0.80-0.84:  0 0 0 0 0
0.85-0.89:  0 0 0 0 2
0.90-0.94:  1 0 0 11 22
0.95-1.00:  38 83 135 110 108 2